package penguin.masking;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.commons.LocalVariablesSorter;
import penguin.masking.asm.CustomMethodVisitor;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.*;

import static org.objectweb.asm.Opcodes.ASM7;

/**
 * @author wuxh
 * @version 1.0.0
 */
public class Interceptors implements ClassFileTransformer {

    List<InterceptorConfig> agents = new ArrayList<>();

    Set<String> classSet = new HashSet<>();//重新加载的类

    public void add(InterceptorConfig interceptor) {
        agents.add(interceptor);
        classSet.add(interceptor.getTransformClassName());
    }

    public Set<String> reloadClz() {
        return classSet;
    }

    public boolean containsClass(String className) {
        return classSet.contains(className);
    }

    public CustomMethodVisitor interceptor(String className, String name, String descriptor) {
        for (InterceptorConfig agent : agents) {
            CustomMethodVisitor visitor = agent.visitor(className, name, descriptor);
            if (visitor == null) {
                continue;
            }
            Logger.debug("InterceptorConfig:" + agent.getAgentMethod() + ":" + agent.getTransformClassName() + ":" + agent.getMethod() + ":" + agent.getDescriptor());
            return visitor;
        }
        return null;
    }

    //原class字节保存，避免重复增强
    private static final Map<String, byte[]> OLD_CLASSBUFFER_MAP = new HashMap<>();

    private boolean restore = false;

    public void setRestore(boolean restore) {
        this.restore = restore;
    }

    public byte[] classfileBuffer(String _className, byte[] classfileBuffer) {
        if (!containsClass(_className)) {
            return classfileBuffer;
        }
        byte[] oldbuffer = OLD_CLASSBUFFER_MAP.get(_className);

        //恢复原来字节码
        if (restore) {
            return oldbuffer == null ? classfileBuffer : oldbuffer;
        }
        if (oldbuffer == null) {
            OLD_CLASSBUFFER_MAP.put(_className, classfileBuffer);
        } else {
            classfileBuffer = oldbuffer;
        }
        Logger.debug("transform: " + _className);
        ClassReader cr = new ClassReader(classfileBuffer);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES);
        ClassVisitor cv = new JdbcClassVisitor(_className, cw);
        cr.accept(cv, ClassReader.SKIP_DEBUG);
        return cw.toByteArray();
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        String _className = className.replaceAll("/", ".");
        return classfileBuffer(_className, classfileBuffer);
    }


    public class JdbcClassVisitor extends ClassVisitor {

        private String className;

        public JdbcClassVisitor(String className, ClassVisitor classVisitor) {
            super(ASM7, classVisitor);
            this.className = className;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            if (mv == null) {
                return null;
            }
            CustomMethodVisitor visitor = interceptor(className, name, descriptor);
            if (visitor != null) {
                LocalVariablesSorter _mv = new LocalVariablesSorter(access, descriptor, visitor);
                visitor.setParentMethodVisitor(mv);
                visitor.setLvs(_mv);
                return visitor;
            }

            return mv;
        }
    }
}
